import timm
from pprint import pprint
from mmcv.cnn import get_model_complexity_info
model_names = timm.list_models(pretrained=True)
# pprint(model_names)
#mobilenetv2_100
#swin_tiny_patch4_window7_224
import timm
import torch
# from ptflops import get_model_complexity_info

m = timm.create_model('resnet101', pretrained=True,features_only=True)
# print(m)

dict_name = list(m.state_dict())
for i, p in enumerate(dict_name):
    print(i, p)
o = m(torch.randn(2, 3, 256, 256))
for x in o:
  print(x.shape)
#torch.Size([2, 32, 28, 28])
# torch.Size([2, 96, 14, 14])
# torch.Size([2, 320, 7, 7])
# macs, params = get_model_complexity_info(m, (3,320,320), as_strings=True,
#                                          print_per_layer_stat=True)
# print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
# print('{:<30}  {:<8}'.format('Number of parameters: ', params))
#选取stage很重要？

# flops, params = get_model_complexity_info(m, (3,224,224))
# print(flops)
# print(params)